{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom components\n",
    "\n",
    "As I mentioned earlier in the example notebooks, and also in the `README`, it is possible to customise almost every component in `pytorch-widedeep`.\n",
    "\n",
    "Let's now go through a couple of simple examples to illustrate how that could be done. \n",
    "\n",
    "First let's load and process the data \"as usual\", let's start with a regression and the [airbnb](http://insideairbnb.com/get-the-data.html) dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/javierrodriguezzaurin/.pyenv/versions/3.10.13/envs/widedeep310/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import torch\n",
    "\n",
    "from torch import Tensor\n",
    "from pytorch_widedeep import Trainer\n",
    "from pytorch_widedeep.preprocessing import (\n",
    "    WidePreprocessor,\n",
    "    TabPreprocessor,\n",
    "    TextPreprocessor,\n",
    "    ImagePreprocessor,\n",
    ")\n",
    "from pytorch_widedeep.models import (\n",
    "    Wide,\n",
    "    TabMlp,\n",
    "    Vision,\n",
    "    BasicRNN,\n",
    "    WideDeep,\n",
    ")\n",
    "from pytorch_widedeep.losses import RMSELoss\n",
    "from pytorch_widedeep.initializers import *\n",
    "from pytorch_widedeep.callbacks import *\n",
    "from pytorch_widedeep.datasets import load_adult"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>host_id</th>\n",
       "      <th>description</th>\n",
       "      <th>host_listings_count</th>\n",
       "      <th>host_identity_verified</th>\n",
       "      <th>neighbourhood_cleansed</th>\n",
       "      <th>latitude</th>\n",
       "      <th>longitude</th>\n",
       "      <th>is_location_exact</th>\n",
       "      <th>property_type</th>\n",
       "      <th>...</th>\n",
       "      <th>amenity_wide_entrance</th>\n",
       "      <th>amenity_wide_entrance_for_guests</th>\n",
       "      <th>amenity_wide_entryway</th>\n",
       "      <th>amenity_wide_hallways</th>\n",
       "      <th>amenity_wifi</th>\n",
       "      <th>amenity_window_guards</th>\n",
       "      <th>amenity_wine_cooler</th>\n",
       "      <th>security_deposit</th>\n",
       "      <th>extra_people</th>\n",
       "      <th>yield</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>13913.jpg</td>\n",
       "      <td>54730</td>\n",
       "      <td>My bright double bedroom with a large window h...</td>\n",
       "      <td>4.0</td>\n",
       "      <td>f</td>\n",
       "      <td>Islington</td>\n",
       "      <td>51.56802</td>\n",
       "      <td>-0.11121</td>\n",
       "      <td>t</td>\n",
       "      <td>apartment</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>100.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>12.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15400.jpg</td>\n",
       "      <td>60302</td>\n",
       "      <td>Lots of windows and light.  St Luke's Gardens ...</td>\n",
       "      <td>1.0</td>\n",
       "      <td>t</td>\n",
       "      <td>Kensington and Chelsea</td>\n",
       "      <td>51.48796</td>\n",
       "      <td>-0.16898</td>\n",
       "      <td>t</td>\n",
       "      <td>apartment</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>150.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>109.50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>17402.jpg</td>\n",
       "      <td>67564</td>\n",
       "      <td>Open from June 2018 after a 3-year break, we a...</td>\n",
       "      <td>19.0</td>\n",
       "      <td>t</td>\n",
       "      <td>Westminster</td>\n",
       "      <td>51.52098</td>\n",
       "      <td>-0.14002</td>\n",
       "      <td>t</td>\n",
       "      <td>apartment</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>350.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>149.65</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>24328.jpg</td>\n",
       "      <td>41759</td>\n",
       "      <td>Artist house, bright high ceiling rooms, priva...</td>\n",
       "      <td>2.0</td>\n",
       "      <td>t</td>\n",
       "      <td>Wandsworth</td>\n",
       "      <td>51.47298</td>\n",
       "      <td>-0.16376</td>\n",
       "      <td>t</td>\n",
       "      <td>other</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>250.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>215.60</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>25023.jpg</td>\n",
       "      <td>102813</td>\n",
       "      <td>Large, all comforts, 2-bed flat; first floor; ...</td>\n",
       "      <td>1.0</td>\n",
       "      <td>f</td>\n",
       "      <td>Wandsworth</td>\n",
       "      <td>51.44687</td>\n",
       "      <td>-0.21874</td>\n",
       "      <td>t</td>\n",
       "      <td>apartment</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>250.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>79.35</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 223 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          id  host_id                                        description  \\\n",
       "0  13913.jpg    54730  My bright double bedroom with a large window h...   \n",
       "1  15400.jpg    60302  Lots of windows and light.  St Luke's Gardens ...   \n",
       "2  17402.jpg    67564  Open from June 2018 after a 3-year break, we a...   \n",
       "3  24328.jpg    41759  Artist house, bright high ceiling rooms, priva...   \n",
       "4  25023.jpg   102813  Large, all comforts, 2-bed flat; first floor; ...   \n",
       "\n",
       "   host_listings_count host_identity_verified  neighbourhood_cleansed  \\\n",
       "0                  4.0                      f               Islington   \n",
       "1                  1.0                      t  Kensington and Chelsea   \n",
       "2                 19.0                      t             Westminster   \n",
       "3                  2.0                      t              Wandsworth   \n",
       "4                  1.0                      f              Wandsworth   \n",
       "\n",
       "   latitude  longitude is_location_exact property_type  ...  \\\n",
       "0  51.56802   -0.11121                 t     apartment  ...   \n",
       "1  51.48796   -0.16898                 t     apartment  ...   \n",
       "2  51.52098   -0.14002                 t     apartment  ...   \n",
       "3  51.47298   -0.16376                 t         other  ...   \n",
       "4  51.44687   -0.21874                 t     apartment  ...   \n",
       "\n",
       "  amenity_wide_entrance  amenity_wide_entrance_for_guests  \\\n",
       "0                     1                                 0   \n",
       "1                     0                                 0   \n",
       "2                     0                                 0   \n",
       "3                     0                                 0   \n",
       "4                     0                                 0   \n",
       "\n",
       "   amenity_wide_entryway  amenity_wide_hallways  amenity_wifi  \\\n",
       "0                      0                      0             1   \n",
       "1                      0                      0             1   \n",
       "2                      0                      0             1   \n",
       "3                      0                      0             1   \n",
       "4                      0                      0             1   \n",
       "\n",
       "   amenity_window_guards  amenity_wine_cooler security_deposit extra_people  \\\n",
       "0                      0                    0            100.0         15.0   \n",
       "1                      0                    0            150.0          0.0   \n",
       "2                      0                    0            350.0         10.0   \n",
       "3                      0                    0            250.0          0.0   \n",
       "4                      0                    0            250.0         11.0   \n",
       "\n",
       "    yield  \n",
       "0   12.00  \n",
       "1  109.50  \n",
       "2  149.65  \n",
       "3  215.60  \n",
       "4   79.35  \n",
       "\n",
       "[5 rows x 223 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"../tmp_data/airbnb/airbnb_sample.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# There are a number of columns that are already binary. Therefore, no need to one hot encode them\n",
    "crossed_cols = [(\"property_type\", \"room_type\")]\n",
    "already_dummies = [c for c in df.columns if \"amenity\" in c] + [\"has_house_rules\"]\n",
    "wide_cols = [\n",
    "    \"is_location_exact\",\n",
    "    \"property_type\",\n",
    "    \"room_type\",\n",
    "    \"host_gender\",\n",
    "    \"instant_bookable\",\n",
    "] + already_dummies\n",
    "\n",
    "cat_embed_cols = [(c, 16) for c in df.columns if \"catg\" in c] + [\n",
    "    (\"neighbourhood_cleansed\", 64),\n",
    "    (\"cancellation_policy\", 16),\n",
    "]\n",
    "continuous_cols = [\"latitude\", \"longitude\", \"security_deposit\", \"extra_people\"]\n",
    "# it does not make sense to standarised Latitude and Longitude\n",
    "already_standard = [\"latitude\", \"longitude\"]\n",
    "\n",
    "# text and image colnames\n",
    "text_col = \"description\"\n",
    "img_col = \"id\"\n",
    "\n",
    "# path to pretrained word embeddings and the images\n",
    "word_vectors_path = \"../tmp_data/glove.6B/glove.6B.100d.txt\"\n",
    "img_path = \"../tmp_data/airbnb/property_picture\"\n",
    "\n",
    "# target\n",
    "target_col = \"yield\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = df[target_col].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/javierrodriguezzaurin/Projects/pytorch-widedeep/pytorch_widedeep/preprocessing/tab_preprocessor.py:358: UserWarning: Continuous columns will not be normalised\n",
      "  warnings.warn(\"Continuous columns will not be normalised\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The vocabulary contains 2192 tokens\n",
      "Indexing word vectors...\n",
      "Loaded 400000 word vectors\n",
      "Preparing embeddings matrix...\n",
      "2175 words in the vocabulary had ../tmp_data/glove.6B/glove.6B.100d.txt vectors and appear more than 5 times\n",
      "Reading Images from ../tmp_data/airbnb/property_picture\n",
      "Resizing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 1001/1001 [00:02<00:00, 497.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing normalisation metrics\n"
     ]
    }
   ],
   "source": [
    "wide_preprocessor = WidePreprocessor(wide_cols=wide_cols, crossed_cols=crossed_cols)\n",
    "X_wide = wide_preprocessor.fit_transform(df)\n",
    "\n",
    "tab_preprocessor = TabPreprocessor(\n",
    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
    ")\n",
    "X_tab = tab_preprocessor.fit_transform(df)\n",
    "\n",
    "text_preprocessor = TextPreprocessor(\n",
    "    word_vectors_path=word_vectors_path, text_col=text_col\n",
    ")\n",
    "X_text = text_preprocessor.fit_transform(df)\n",
    "\n",
    "image_processor = ImagePreprocessor(img_col=img_col, img_path=img_path)\n",
    "X_images = image_processor.fit_transform(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to build a wide and deep model. Three of the four components we will use are included in this package, and they will be combined with a custom `deeptext` component. Then the fit process will run with a custom loss function.\n",
    "\n",
    "Let's have a look"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Linear model\n",
    "wide = Wide(input_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
    "\n",
    "# DeepDense: 2 Dense layers\n",
    "tab_mlp = TabMlp(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    cat_embed_dropout=0.1,\n",
    "    continuous_cols=continuous_cols,\n",
    "    mlp_hidden_dims=[128, 64],\n",
    "    mlp_dropout=0.1,\n",
    ")\n",
    "\n",
    "# Pretrained Resnet 18\n",
    "resnet = Vision(pretrained_model_name=\"resnet18\", n_trainable=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom `deeptext`\n",
    "\n",
    "Standard Pytorch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDeepText(nn.Module):\n",
    "    def __init__(self, vocab_size, padding_idx=1, embed_dim=100, hidden_dim=64):\n",
    "        super(MyDeepText, self).__init__()\n",
    "\n",
    "        # word/token embeddings\n",
    "        self.word_embed = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)\n",
    "\n",
    "        # stack of RNNs\n",
    "        self.rnn = nn.GRU(\n",
    "            embed_dim,\n",
    "            hidden_dim,\n",
    "            num_layers=2,\n",
    "            bidirectional=True,\n",
    "            batch_first=True,\n",
    "        )\n",
    "\n",
    "        # Remember, this MUST be defined. If not WideDeep will through an error\n",
    "        self.output_dim = hidden_dim * 2\n",
    "\n",
    "    def forward(self, X):\n",
    "        embed = self.word_embed(X.long())\n",
    "        o, h = self.rnn(embed)\n",
    "        return torch.cat((h[-2], h[-1]), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "mydeeptext = MyDeepText(vocab_size=len(text_preprocessor.vocab.itos))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = WideDeep(wide=wide, deeptabular=tab_mlp, deeptext=mydeeptext, deepimage=resnet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom loss function\n",
    "\n",
    "Loss functions must simply inherit pytorch's `nn.Module`. For example, let's say we want to use `RMSE` (note that this is already available in the package, but I will pass it here as a custom loss for illustration purposes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSELoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        \"\"\"root mean squared error\"\"\"\n",
    "        super().__init__()\n",
    "        self.mse = nn.MSELoss()\n",
    "\n",
    "    def forward(self, input: Tensor, target: Tensor) -> Tensor:\n",
    "        return torch.sqrt(self.mse(input, target))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now we just instantiate the ``Trainer`` as usual. Needless to say, but this runs with 1000 random observations, so loss and metric values are meaningless. This is just an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(model, objective=\"regression\", custom_loss_function=RMSELoss())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|███████████████████████████████████████████████████████████████████████████████████████| 25/25 [00:23<00:00,  1.07it/s, loss=126]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:05<00:00,  1.24it/s, loss=97.4]\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(\n",
    "    X_wide=X_wide,\n",
    "    X_tab=X_tab,\n",
    "    X_text=X_text,\n",
    "    X_img=X_images,\n",
    "    target=target,\n",
    "    n_epochs=1,\n",
    "    batch_size=32,\n",
    "    val_split=0.2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to model components and loss functions, we can also use custom callbacks or custom metrics. The former need to be of type `Callback` and the latter need to be of type `Metric`. See:\n",
    "\n",
    "```python\n",
    "pytorch-widedeep.callbacks\n",
    "```\n",
    "and \n",
    "\n",
    "```python\n",
    "pytorch-widedeep.metrics\n",
    "```\n",
    "\n",
    "For this example let me use the adult dataset. Again, we first prepare the data as usual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>educational-num</th>\n",
       "      <th>marital-status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>native-country</th>\n",
       "      <th>income</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>25</td>\n",
       "      <td>Private</td>\n",
       "      <td>226802</td>\n",
       "      <td>11th</td>\n",
       "      <td>7</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>89814</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Farming-fishing</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>336951</td>\n",
       "      <td>Assoc-acdm</td>\n",
       "      <td>12</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Protective-serv</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&gt;50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>44</td>\n",
       "      <td>Private</td>\n",
       "      <td>160323</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>10</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Husband</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>7688</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&gt;50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>18</td>\n",
       "      <td>?</td>\n",
       "      <td>103497</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>10</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>?</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>White</td>\n",
       "      <td>Female</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>30</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass  fnlwgt     education  educational-num      marital-status  \\\n",
       "0   25    Private  226802          11th                7       Never-married   \n",
       "1   38    Private   89814       HS-grad                9  Married-civ-spouse   \n",
       "2   28  Local-gov  336951    Assoc-acdm               12  Married-civ-spouse   \n",
       "3   44    Private  160323  Some-college               10  Married-civ-spouse   \n",
       "4   18          ?  103497  Some-college               10       Never-married   \n",
       "\n",
       "          occupation relationship   race  gender  capital-gain  capital-loss  \\\n",
       "0  Machine-op-inspct    Own-child  Black    Male             0             0   \n",
       "1    Farming-fishing      Husband  White    Male             0             0   \n",
       "2    Protective-serv      Husband  White    Male             0             0   \n",
       "3  Machine-op-inspct      Husband  Black    Male          7688             0   \n",
       "4                  ?    Own-child  White  Female             0             0   \n",
       "\n",
       "   hours-per-week native-country income  \n",
       "0              40  United-States  <=50K  \n",
       "1              50  United-States  <=50K  \n",
       "2              40  United-States   >50K  \n",
       "3              40  United-States   >50K  \n",
       "4              30  United-States  <=50K  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = load_adult(as_frame=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>educational_num</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>25</td>\n",
       "      <td>Private</td>\n",
       "      <td>226802</td>\n",
       "      <td>11th</td>\n",
       "      <td>7</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>89814</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Farming-fishing</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>336951</td>\n",
       "      <td>Assoc-acdm</td>\n",
       "      <td>12</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Protective-serv</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>44</td>\n",
       "      <td>Private</td>\n",
       "      <td>160323</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>10</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Husband</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>7688</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>18</td>\n",
       "      <td>?</td>\n",
       "      <td>103497</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>10</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>?</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>White</td>\n",
       "      <td>Female</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>30</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass  fnlwgt     education  educational_num      marital_status  \\\n",
       "0   25    Private  226802          11th                7       Never-married   \n",
       "1   38    Private   89814       HS-grad                9  Married-civ-spouse   \n",
       "2   28  Local-gov  336951    Assoc-acdm               12  Married-civ-spouse   \n",
       "3   44    Private  160323  Some-college               10  Married-civ-spouse   \n",
       "4   18          ?  103497  Some-college               10       Never-married   \n",
       "\n",
       "          occupation relationship   race  gender  capital_gain  capital_loss  \\\n",
       "0  Machine-op-inspct    Own-child  Black    Male             0             0   \n",
       "1    Farming-fishing      Husband  White    Male             0             0   \n",
       "2    Protective-serv      Husband  White    Male             0             0   \n",
       "3  Machine-op-inspct      Husband  Black    Male          7688             0   \n",
       "4                  ?    Own-child  White  Female             0             0   \n",
       "\n",
       "   hours_per_week native_country  income_label  \n",
       "0              40  United-States             0  \n",
       "1              50  United-States             0  \n",
       "2              40  United-States             1  \n",
       "3              40  United-States             1  \n",
       "4              30  United-States             0  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For convenience, we'll replace '-' with '_'\n",
    "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
    "# binary target\n",
    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
    "df.drop(\"income\", axis=1, inplace=True)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define wide, crossed and deep tabular columns\n",
    "wide_cols = [\n",
    "    \"workclass\",\n",
    "    \"education\",\n",
    "    \"marital_status\",\n",
    "    \"occupation\",\n",
    "    \"relationship\",\n",
    "    \"race\",\n",
    "    \"gender\",\n",
    "    \"native_country\",\n",
    "]\n",
    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]\n",
    "cat_embed_cols = [\n",
    "    \"workclass\",\n",
    "    \"education\",\n",
    "    \"marital_status\",\n",
    "    \"occupation\",\n",
    "    \"relationship\",\n",
    "    \"race\",\n",
    "    \"gender\",\n",
    "    \"capital_gain\",\n",
    "    \"capital_loss\",\n",
    "    \"native_country\",\n",
    "]\n",
    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
    "target_col = \"income_label\"\n",
    "target = df[target_col].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/javierrodriguezzaurin/Projects/pytorch-widedeep/pytorch_widedeep/preprocessing/tab_preprocessor.py:358: UserWarning: Continuous columns will not be normalised\n",
      "  warnings.warn(\"Continuous columns will not be normalised\")\n"
     ]
    }
   ],
   "source": [
    "# wide\n",
    "wide_preprocessor = WidePreprocessor(wide_cols=wide_cols, crossed_cols=crossed_cols)\n",
    "X_wide = wide_preprocessor.fit_transform(df)\n",
    "\n",
    "# deeptabular\n",
    "tab_preprocessor = TabPreprocessor(\n",
    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
    ")\n",
    "X_tab = tab_preprocessor.fit_transform(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "wide = Wide(input_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
    "tab_mlp = TabMlp(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    continuous_cols=continuous_cols,\n",
    "    mlp_hidden_dims=[128, 64],\n",
    "    mlp_dropout=0.2,\n",
    "    mlp_activation=\"leaky_relu\",\n",
    ")\n",
    "model = WideDeep(wide=wide, deeptabular=tab_mlp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom metric\n",
    "\n",
    "Let's say we want to use our own accuracy metric (again, this is already available in the package, but I will pass it here as a custom loss for illustration purposes). \n",
    "\n",
    "This could be done as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_widedeep.metrics import Metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Accuracy(Metric):\n",
    "    def __init__(self, top_k: int = 1):\n",
    "        super(Accuracy, self).__init__()\n",
    "\n",
    "        self.top_k = top_k\n",
    "        self.correct_count = 0\n",
    "        self.total_count = 0\n",
    "\n",
    "        #  metric name needs to be defined\n",
    "        self._name = \"acc\"\n",
    "\n",
    "    def reset(self):\n",
    "        self.correct_count = 0\n",
    "        self.total_count = 0\n",
    "\n",
    "    def __call__(self, y_pred: Tensor, y_true: Tensor) -> np.ndarray:\n",
    "        num_classes = y_pred.size(1)\n",
    "\n",
    "        if num_classes == 1:\n",
    "            y_pred = y_pred.round()\n",
    "            y_true = y_true\n",
    "        elif num_classes > 1:\n",
    "            y_pred = y_pred.topk(self.top_k, 1)[1]\n",
    "            y_true = y_true.view(-1, 1).expand_as(y_pred)\n",
    "\n",
    "        self.correct_count += y_pred.eq(y_true).sum().item()\n",
    "        self.total_count += len(y_pred)\n",
    "        accuracy = float(self.correct_count) / float(self.total_count)\n",
    "        return np.array(accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Custom Callback\n",
    "\n",
    "Let's code a callback that records the current epoch at the beginning and the end of each epoch (silly, but you know, this is just an example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# have a look to the class\n",
    "from pytorch_widedeep.callbacks import Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SillyCallback(Callback):\n",
    "    def on_train_begin(self, logs=None):\n",
    "        # recordings will be the trainer object attributes\n",
    "        self.trainer.silly_callback = {}\n",
    "\n",
    "        self.trainer.silly_callback[\"beginning\"] = []\n",
    "        self.trainer.silly_callback[\"end\"] = []\n",
    "\n",
    "    def on_epoch_begin(self, epoch, logs=None):\n",
    "        self.trainer.silly_callback[\"beginning\"].append(epoch + 1)\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None, metric=None):\n",
    "        self.trainer.silly_callback[\"end\"].append(epoch + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now, as usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model, objective=\"binary\", metrics=[Accuracy], callbacks=[SillyCallback]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|███████████████████████████████████████████████████████████| 611/611 [00:06<00:00, 94.39it/s, loss=0.411, metrics={'acc': 0.814}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 121.91it/s, loss=0.327, metrics={'acc': 0.8449}]\n",
      "epoch 2: 100%|██████████████████████████████████████████████████████████| 611/611 [00:07<00:00, 85.39it/s, loss=0.324, metrics={'acc': 0.8495}]\n",
      "valid: 100%|████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 88.68it/s, loss=0.298, metrics={'acc': 0.8612}]\n",
      "epoch 3: 100%|██████████████████████████████████████████████████████████| 611/611 [00:08<00:00, 74.35it/s, loss=0.302, metrics={'acc': 0.8593}]\n",
      "valid: 100%|████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 100.51it/s, loss=0.29, metrics={'acc': 0.8665}]\n",
      "epoch 4: 100%|██████████████████████████████████████████████████████████| 611/611 [00:08<00:00, 73.83it/s, loss=0.292, metrics={'acc': 0.8637}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 105.98it/s, loss=0.286, metrics={'acc': 0.8695}]\n",
      "epoch 5: 100%|███████████████████████████████████████████████████████████| 611/611 [00:08<00:00, 72.15it/s, loss=0.286, metrics={'acc': 0.866}]\n",
      "valid: 100%|████████████████████████████████████████████████████████████| 153/153 [00:01<00:00, 92.27it/s, loss=0.284, metrics={'acc': 0.8698}]\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(\n",
    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'beginning': [1, 2, 3, 4, 5], 'end': [1, 2, 3, 4, 5]}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.silly_callback"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
